/**
* \brief 
* \author pengcheng (pengcheng@yslrpch@126.com)
* \date 2020-05-30
* \attention CopyrightÃ‚Â©ADC Technology(tianjin)Co.Ltd
* \attention Refer to COPYRIGHT.txt for complete terms of copyright notice
*/
#include <detection_vision/tensorrt/image_processor.h>

namespace tensorrt_inference
{


bool ImageProcessor::LoadImageAsNCHW(std::vector<float>& output,
                                     const bool in_RGB_order)
{
	const size_t channel_size = 3;
	if ((image_info_.mean.size() != 3) || (image_info_.dev.size() != 3))
	{
		LOG(ERROR) << "Size of mean/dev should be 3";
		return false;
	}
    // reserve for output
    output.clear();
    output.resize(image_info_.paths.size() * channel_size * image_info_.height * image_info_.width);
    size_t idx = 0;
	// load raw image
    for (size_t i = 0; i < image_info_.paths.size(); ++i)
    {
        cv::Mat img_raw = cv::imread(image_info_.paths[i]);
        if (img_raw.data == nullptr)
        {
            LOG(ERROR) << "Failed to load image from " << image_info_.paths[i];
            return false;
        }
        const size_t in_height = img_raw.rows;
        const size_t in_width = img_raw.cols;
        if ((in_height > image_info_.height) || (in_width > image_info_.width))
        {
            LOG(ERROR) << "Input size is larger than output size, can not handle this case";
            return false;
        }

        // convert from BGR to RGB
        if (in_RGB_order)
        {
            cv::cvtColor(img_raw, img_raw, cv::COLOR_BGR2RGB);
        }

        // calculate padding size
        const size_t top_padding_size = (image_info_.height - in_height) / 2;
        const size_t left_padding_size = (image_info_.width - in_width) / 2;

        // pad image
        const uint8_t* in_data = img_raw.ptr<uint8_t>(0);
        for (size_t c = 0; c < channel_size; ++c)
        {
            for (size_t h = 0; h < image_info_.height; ++h)
            {
                for (size_t w = 0; w < image_info_.width; ++w)
                {
                    if ((h >= top_padding_size) &&
                        (h < (in_height + top_padding_size)) &&
                        (w >= left_padding_size) &&
                        (w < (in_width + left_padding_size)))
                    {
                        const size_t tmp_idx = (h - top_padding_size) * channel_size * in_width +
                                               (w - left_padding_size) * channel_size +
                                               c;
                        output[idx++] = float((float(in_data[tmp_idx]) * image_info_.scale -
                                                         image_info_.mean[c]) / image_info_.dev[c]);
                    }
                    else
                    {
                        output[idx++] = float(image_info_.padding_with);
                    }
                }
            }
        }
    }

	return true;
}

}